package org.example.yolodemo.util;

import ai.onnxruntime.*;
import com.alibaba.fastjson.JSONObject;
import org.opencv.core.*;
import org.opencv.imgproc.Imgproc;
import org.opencv.videoio.VideoCapture;
import org.opencv.videoio.VideoWriter;
import nu.pattern.OpenCV;

import java.io.File;
import java.io.IOException;
import java.net.URL;
import java.net.URLDecoder;
import java.nio.FloatBuffer;
import java.text.DecimalFormat;
import java.util.*;
import java.util.List;

import static org.opencv.videoio.Videoio.*;

/**
 * ONNX模型加载与目标检测工具类（扩展视频处理功能）
 */
public class VideosUtil {
    static {
        try {
            // 使用nu.pattern提供的OpenCV加载器
            OpenCV.loadLocally();
            System.out.println("OpenCV native library loaded successfully");
            
            // 初始化ONNX运行时环境
            env = OrtEnvironment.getEnvironment();
            System.out.println("ONNX Runtime environment initialized successfully");
        } catch (Exception e) {
            System.err.println("Error during initialization: " + e.getMessage());
            e.printStackTrace();
        }
    }

    private static OrtEnvironment env;
    private static OrtSession session;
    private static JSONObject names;
    private static long count;
    private static long channels;
    private static long netHeight;
    private static long netWidth;
    private static float srcw;
    private static float srch;
    private static float confThreshold = 0.25f;
    private static float nmsThreshold = 0.5f;
    private static Mat src;

    // 视频处理相关参数
    private static final DecimalFormat df = new DecimalFormat("0000"); // 帧编号格式化

    /**
     * 加载ONNX模型及相关配置
     * @param path ONNX模型文件路径
     */
    public static void load(String path) {
        try {
            if (env == null) {
                env = OrtEnvironment.getEnvironment();
                System.out.println("ONNX Runtime environment initialized in load()");
            }

            // 如果已经有会话，先关闭它
            if (session != null) {
                session.close();
            }

            // 创建新的会话
            session = env.createSession(path, new OrtSession.SessionOptions());
            System.out.println("ONNX model session created successfully");

            // 获取模型元数据
            OnnxModelMetadata metadata = session.getMetadata();
            String nameClass = metadata.getCustomMetadata().get("names");
            names = JSONObject.parseObject(nameClass.replace("\"", "\"\""));

            // 获取输入信息
            Map<String, NodeInfo> infoMap = session.getInputInfo();
            TensorInfo nodeInfo = (TensorInfo) infoMap.get("images").getInfo();
            count = nodeInfo.getShape()[0];
            channels = nodeInfo.getShape()[1];
            netHeight = nodeInfo.getShape()[2];
            netWidth = nodeInfo.getShape()[3];

            System.out.println("Model loaded successfully with input shape: " + 
                             count + "x" + channels + "x" + netHeight + "x" + netWidth);
        } catch (Exception e) {
            System.err.println("Error loading model: " + e.getMessage());
            e.printStackTrace();
            throw new RuntimeException("Failed to load ONNX model", e);
        }
    }

    /**
     * 对指定路径的视频进行目标检测
     * @param inputVideoPath 输入视频路径
     * @param outputVideoPath 输出视频路径（带识别结果）
     * @return 处理状态信息
     * @throws Exception 可能抛出的异常
     */
    public static Map<String, Object> processVideo(String inputVideoPath, String outputVideoPath) throws Exception {
        // 确保模型已加载
        if (session == null) {
            throw new IllegalStateException("ONNX model not loaded. Please call load() first.");
        }

        Map<String, Object> result = new HashMap<>();
        result.put("success", false);
        result.put("message", "处理失败");

        // 打开视频文件
        VideoCapture capture = new VideoCapture(inputVideoPath);
        if (!capture.isOpened()) {
            result.put("message", "无法打开视频文件");
            return result;
        }

        // 获取视频属性
        double fps = capture.get(CAP_PROP_FPS);
        int width = (int) capture.get(CAP_PROP_FRAME_WIDTH);
        int height = (int) capture.get(CAP_PROP_FRAME_HEIGHT);
        long frameCount = (long) capture.get(CAP_PROP_FRAME_COUNT);

        // 创建视频写入器
        VideoWriter writer = new VideoWriter();
        boolean isMp4 = outputVideoPath.endsWith(".mp4");
        int fourcc = isMp4 ? VideoWriter.fourcc('m', 'p', '4', 'v') : VideoWriter.fourcc('X', 'V', 'I', 'D');

        if (!writer.open(outputVideoPath, fourcc, fps, new Size(width, height), true)) {
            capture.release();
            result.put("message", "无法创建输出视频文件");
            return result;
        }

        Mat frame = new Mat();
        long processedFrames = 0;
        long startTime = System.currentTimeMillis();

        try {
            // 逐帧处理视频
            while (capture.read(frame)) {
                if (frame.empty()) {
                    break;
                }

                // 处理当前帧
                Mat processedFrame = processFrame(frame);

                // 写入处理后的帧
                writer.write(processedFrame);

                processedFrames++;
                if (processedFrames % 10 == 0) {
                    System.out.println("已处理帧: " + processedFrames + "/" + frameCount);
                }
            }

            result.put("success", true);
            result.put("message", "视频处理完成");
            result.put("processedFrames", processedFrames);
            result.put("frameCount", frameCount);
            result.put("processingTimeMs", System.currentTimeMillis() - startTime);

        } catch (Exception e) {
            e.printStackTrace();
            result.put("message", "处理过程中发生错误: " + e.getMessage());
        } finally {
            // 释放资源
            capture.release();
            writer.release();
            if (frame != null) frame.release();
        }

        return result;
    }

    /**
     * 处理单帧图像，添加目标检测结果
     * @param frame 输入帧
     * @return 处理后的帧
     * @throws Exception 可能抛出的异常
     */
    private static Mat processFrame(Mat frame) throws Exception {
        if (frame.empty()) {
            throw new IllegalArgumentException("Input frame is empty");
        }

        // 保存原始尺寸
        srcw = frame.width();
        srch = frame.height();

        // 创建输入图像的副本
        src = frame.clone();

        // 确保目标尺寸大于0
        if (netWidth <= 0 || netHeight <= 0) {
            netWidth = 640;  // 设置默认值
            netHeight = 640;
        }

        // 计算缩放比例
        double scale = Math.min((double)netWidth/srcw, (double)netHeight/srch);
        int newWidth = (int)(srcw * scale);
        int newHeight = (int)(srch * scale);

        // 调整图像大小
        Mat resizedFrame = new Mat();
        Imgproc.resize(src, resizedFrame, new Size(newWidth, newHeight));

        // 创建方形画布
        Mat squareFrame = new Mat((int) netHeight, (int) netWidth, resizedFrame.type(), new Scalar(0, 0, 0));
        
        // 将调整大小的图像复制到方形画布的中心
        int x = (int) ((netWidth - newWidth) / 2);
        int y = (int) ((netHeight - newHeight) / 2);
        resizedFrame.copyTo(squareFrame.submat(new Rect(x, y, newWidth, newHeight)));

        // 执行目标检测
        Map<Object, Object> detectionResult = predictor();
        List<ArrayList<Float>> boxes = (List<ArrayList<Float>>) detectionResult.get("boxes");
        JSONObject classNames = (JSONObject) detectionResult.get("classNames");

        // 在原始帧上绘制检测结果
        for (ArrayList<Float> box : boxes) {
            float x1 = box.get(0);
            float y1 = box.get(1);
            float x2 = box.get(2);
            float y2 = box.get(3);
            float confidence = box.get(4);
            int classNameIndex = box.get(5).intValue();
            String className = classNames.getString(String.valueOf(classNameIndex));

            // 绘制矩形框
            Imgproc.rectangle(frame, new Point(x1, y1), new Point(x2, y2), new Scalar(0, 0, 255), 2);

            // 绘制类别名称和置信度
            String label = className + ": " + new DecimalFormat("#.##").format(confidence);
            Imgproc.putText(frame, label, new Point(x1, y1 - 5), Imgproc.FONT_HERSHEY_SIMPLEX, 0.7, new Scalar(255, 0, 0), 2);
        }

        // 释放资源
        resizedFrame.release();
        squareFrame.release();

        return frame;
    }

    public static OnnxTensor transferTensor(Mat dst) {
        try {
            if (env == null) {
                env = OrtEnvironment.getEnvironment();
                System.out.println("ONNX Runtime environment initialized in transferTensor()");
            }

            Imgproc.cvtColor(dst, dst, Imgproc.COLOR_BGR2RGB);
            dst.convertTo(dst, CvType.CV_32FC1, 1. / 255);
            float[] whc = new float[Long.valueOf(channels).intValue() * Long.valueOf(netWidth).intValue() * Long.valueOf(netHeight).intValue()];
            dst.get(0, 0, whc);
            float[] chw = whc2cwh(whc);

            return OnnxTensor.createTensor(env, FloatBuffer.wrap(chw), new long[]{count, channels, netWidth, netHeight});
        } catch (Exception e) {
            System.err.println("Error creating ONNX tensor: " + e.getMessage());
            e.printStackTrace();
            throw new RuntimeException("Failed to create ONNX tensor", e);
        }
    }

    public static float[] whc2cwh(float[] src) {
        float[] chw = new float[src.length];
        int j = 0;
        for (int ch = 0; ch < 3; ++ch) {
            for (int i = ch; i < src.length; i += 3) {
                chw[j] = src[i];
                j++;
            }
        }
        return chw;
    }

    public static Map<Object, Object> predictor() throws Exception {
        // 确保src不为空
        if (src == null || src.empty()) {
            throw new IllegalStateException("Source image is null or empty");
        }

        float scaleW = srcw / netWidth;
        float scaleH = srch / netHeight;

        // 调整图像大小
        Mat dst = new Mat();
        Imgproc.resize(src, dst, new Size(netWidth, netHeight));
        
        OnnxTensor tensor = transferTensor(dst);
        OrtSession.Result result = session.run(Collections.singletonMap("images", tensor));
        
        // 释放资源
        dst.release();
        
        OnnxTensor res = (OnnxTensor) result.get(0);
        float[][][] dataRes = (float[][][]) res.getValue();
        float[][] data = dataRes[0];

        float rawData[][] = new float[data[0].length][6];
        for (int i = 0; i < 4; i++) {
            for (int j = 0; j < data[0].length; j++) {
                rawData[j][i] = data[i][j];
            }
        }

        for (int i = 0; i < data[0].length; i++) {
            for (int j = 4; j < data.length; j++) {
                if (rawData[i][4] < data[j][i]) {
                    rawData[i][4] = data[j][i];
                    rawData[i][5] = j - 4;
                }
            }
        }

        List<ArrayList<Float>> boxes = new LinkedList<ArrayList<Float>>();
        ArrayList<Float> box = null;
        for (float[] d : rawData) {
            if (d[4] > confThreshold) {
                d[0] = d[0] - d[2] / 2;
                d[1] = d[1] - d[3] / 2;
                d[2] = d[0] + d[2];
                d[3] = d[1] + d[3];
                box = new ArrayList<Float>();
                for (float num : d) {
                    box.add(num);
                }
                if (boxes.size() == 0) {
                    boxes.add(box);
                } else {
                    int i;
                    for (i = 0; i < boxes.size(); i++) {
                        if (box.get(4) > boxes.get(i).get(4)) {
                            boxes.add(i, box);
                            break;
                        }
                    }
                    if (i == boxes.size()) {
                        boxes.add(box);
                    }
                }
            }
        }

        int[] indexs = new int[boxes.size()];
        Arrays.fill(indexs, 1);
        for (int cur = 0; cur < boxes.size(); cur++) {
            if (indexs[cur] == 0) {
                continue;
            }
            ArrayList<Float> curMaxConf = boxes.get(cur);
            for (int i = cur + 1; i < boxes.size(); i++) {
                if (indexs[i] == 0) {
                    continue;
                }
                float classIndex = boxes.get(i).get(5);
                if (classIndex == curMaxConf.get(5)) {
                    float x1 = curMaxConf.get(0);
                    float y1 = curMaxConf.get(1);
                    float x2 = curMaxConf.get(2);
                    float y2 = curMaxConf.get(3);
                    float x3 = boxes.get(i).get(0);
                    float y3 = boxes.get(i).get(1);
                    float x4 = boxes.get(i).get(2);
                    float y4 = boxes.get(i).get(3);
                    if (x1 > x4 || x2 < x3 || y1 > y4 || y2 < y3) {
                        continue;
                    }
                    float intersectionWidth = Math.max(x1, x3) - Math.min(x2, x4);
                    float intersectionHeight = Math.max(y1, y3) - Math.min(y2, y4);
                    float intersectionArea = Math.max(0, intersectionWidth * intersectionHeight);
                    float unionArea = (x2 - x1) * (y2 - y1) + (x4 - x3) * (y4 - y3) - intersectionArea;
                    float iou = intersectionArea / unionArea;
                    indexs[i] = iou > nmsThreshold? 0 : 1;
                }
            }
        }

        List<ArrayList<Float>> resBoxes = new LinkedList<ArrayList<Float>>();
        for (int index = 0; index < indexs.length; index++) {
            if (indexs[index] == 1) {
                resBoxes.add(boxes.get(index));
            }
        }
        boxes = resBoxes;

        for (ArrayList<Float> box1 : boxes) {
            box1.set(0, box1.get(0) * scaleW);
            box1.set(1, box1.get(1) * scaleH);
            box1.set(2, box1.get(2) * scaleW);
            box1.set(3, box1.get(3) * scaleH);
        }

        Map<Object, Object> map = new HashMap<Object, Object>();
        map.put("boxes", boxes);
        map.put("classNames", names);
        return map;
    }

    public static String getModelPath(String resourcePath) throws IOException {
        URL resourceUrl = ClassLoader.getSystemResource(resourcePath);
        if (resourceUrl == null) {
            throw new IOException("资源未找到: " + resourcePath);
        }
        String path = URLDecoder.decode(resourceUrl.getPath(), "UTF-8");
        if (System.getProperty("os.name").toLowerCase().contains("win") && path.startsWith("/")) {
            path = path.substring(1);
        }
        File file = new File(path);
        if (!file.exists()) {
            throw new IOException("文件不存在: " + path);
        }
        return path;
    }

    // 测试方法
    public static void main(String[] args) throws Exception {
        String modelPath = getModelPath("models/best.onnx");
        load(modelPath);

        String inputVideo = "C:\\Users\\HP\\IdeaProjects\\YoloDemo\\videos\\sp1.mp4";
        String outputVideo = "C:\\Users\\HP\\IdeaProjects\\YoloDemo\\processed\\result.mp4";

        Map<String, Object> processResult = processVideo(inputVideo, outputVideo);
        System.out.println("处理结果: " + processResult);
    }
}